Added support of 2-3 bit symmetric mode to LoRA-QAT by andreyanufr · Pull Request #4158 · openvinotoolkit/nncf

andreyanufr · 2026-07-30T13:42:29Z

Changes

LoRA-QAT now support two new modes: CompressWeightsMode.INT3_SYM and CompressWeightsMode.INT2_SYM.
New function repack_weights was added to repack 4-8 bit weights to lower bit representation if possible.

Reason for changes

Support of 2-3 bit in compression in QAT and pytorch backend.

Related tickets

CVS-180191

Tests

In progress.

Weight compression - success

… nncf side.

2) Removed 6bit from repacking.

…aining

2) Added optional possibility to save torch checkpoint with dequantized weights.

2) Added README and requirements.txt

2) Removed debug prints.

Copilot

Copilot review overview

🟡 Changes recommended

Unresolved critical and moderate issues remain in low-bit stripping, OpenVINO repacking, and the example workflow.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 2 · 2 · 2

Open (6)

Identify zero-point constants by graph role, not shape · New
Add 2/3-bit DQ packing support or reject validation · New
Include gate_proj in producer search · New
Handle meta-device offloaded weights during equalization · New
Use logger instead of unconditional print calls · New
Strengthen repacking test with graph/output validation · New

What changed in this PR

Adds INT2/INT3 symmetric LoRA-QAT support and OpenVINO weight repacking.

Changes:

Adds 2/3-bit quantization mappings and strip validation.
Implements and exports repack_weights.
Adds tests, models, and a distillation QAT example with documentation.

File	Summary and review notes
`tests/openvino/native/test_model_repacking.py`	Adds repacking coverage; assertions should validate graph correctness and numerical equivalence (nit, 2 votes).
`tests/openvino/native/models.py`	Adds a repacking test model.
`src/nncf/torch/function_hook/strip.py`	2/3-bit DQ still uses 4-bit decompression and packing (critical, 3 votes).
`src/nncf/quantization/quantize_model.py`	Exposes the repacking API.
`src/nncf/quantization/algorithms/weight_compression/torch_backend.py`	Adds INT2/INT3 LoRA quantization mappings.
`src/nncf/quantization/algorithms/weight_compression/openvino_backend.py`	Shape heuristic skips valid signed compressed weights, including 1-D tensors (critical, 2 votes).
`src/nncf/quantization/algorithms/weight_compression/algorithm.py`	Coordinates repacking; direct `print` calls should use the project logger (nit, 4 votes).
`src/nncf/quantization/__init__.py`	Exports `repack_weights`.
`src/nncf/openvino/quantization/quantize_model.py`	Provides the OpenVINO repacking implementation.
`src/nncf/__init__.py`	Exports the public API.
`examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/requirements.txt`	Defines example dependencies.
`examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/README.md`	Documents the example workflow.
`examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit/main.py`	Omits `gate_proj` equalization (moderate, 2 votes), skips meta-device weights (moderate, 3 votes), drops partial accumulation gradients (moderate, 1 vote), and does not return its documented metric (nit, 1 vote).
`.ci/cspell_dict.txt`	Adds example vocabulary.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+    if w.device.type != "meta":
+        return w.data
+    return None


+    repacked_model = nncf.repack_weights(model)
+
+    element_types = get_constant_element_types(repacked_model)
+    assert ov.Type.u3 in element_types, "Expected u3 constant after repacking"
+    assert ov.Type.u2 in element_types, "Expected u2 constant after repacking"


…aining

2) Added arg for deterministic cuda execution. This arg is optional because it can slow down the training.

AlexanderDokuchaev · 2026-09-29T07:16:46Z

+### 1. Create and activate a virtual environment
+
+```bash
+cd examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit


Suggested change

cd examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit

Please keep same instruction steps in examples.

AlexanderDokuchaev · 2026-09-29T07:21:12Z

+### 2. Install the example dependencies from this folder
+
+```bash
+python -m pip install -r requirements.txt
+```
+
+### 3. Install NNCF and other dependencies
+
+```bash
+python3 -m pip install ../../../../ -r requirements.txt
+```


Suggested change

### 2. Install the example dependencies from this folder

```bash

python -m pip install -r requirements.txt

```

### 3. Install NNCF and other dependencies

```bash

python3 -m pip install ../../../../ -r requirements.txt

```

### 2. Install NNCF and other dependencies

```bash

python3 -m pip install ../../../../ -r requirements.txt

?

What about local requirements like datasets ? We do not have them in ../../../../ -r requirements.txt

../../../../ - installing NNCF with pakage's requarements
-r requirements.txt - installing example's requarements

AlexanderDokuchaev · 2026-09-29T07:24:44Z

+cd examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit
+python main.py --pretrained meta-llama/Llama-3.2-1B-Instruct --bits 3 --dataset pile-10k --epochs 1 --output_dir output


Suggested change

cd examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit

python main.py --pretrained meta-llama/Llama-3.2-1B-Instruct --bits 3 --dataset pile-10k --epochs 1 --output_dir output

python main.py

For example, just to run main.py with default arguemts

AlexanderDokuchaev · 2026-09-29T07:27:42Z

+- `--scale_epochs`: extra epochs where only quantizer scales are trained
+- `--linear_lr_scheduler`: use a linear learning-rate decay schedule
+- `--batch_size`: accumulation target
+- `--microbatch_size`: per-step microbatch size


Suggested change

- `--microbatch_size`: per-step microbatch size

- `--microbatch_size`: per-step microbatch size

- `--full_determinism`: to enable deterministic QAT

AlexanderDokuchaev · 2026-09-29T07:29:42Z

+- OpenVINO export is always performed at the end of the run, after checkpoint restoration and stripping.
+- `--resume` reuses the checkpoint if present; otherwise the script initializes from scratch.
+
+For more background on absorbable LoRA and low-bit training-time compression, see the project documentation for QAT LoRA usage and NNCF compression flows.


I'm not sure if this is necessary,
or add links

AlexanderDokuchaev · 2026-09-29T10:14:14Z

        if hook_module._qspec.half_range or hook_module._qspec.narrow_range:
            msg += "Unexpected parameters of quantizers on strip: half_range and narrow_range should be False.\n"
-        if hook_module.num_bits not in [4, 8]:
+        if hook_module.num_bits not in [2, 3, 4, 8]:


Looks this function should work for any num_bits, and this condition can be removed?

AlexanderDokuchaev · 2026-09-29T11:20:33Z

+        supported_src_types = {ov.Type.u4: 4, ov.Type.i4: 4, ov.Type.u8: 8, ov.Type.i8: 8}
+        if ov_node.tensor_view.element_type not in supported_src_types:
+            return False, 0, 0
+        bits = supported_src_types[ov_node.tensor_view.element_type]
+        asym = ov_node.tensor_view.element_type in [ov.Type.u8, ov.Type.u4]
+
+        # only repack for sym types are supported for now: i3, i2
+        if asym:
+            return False, 0, 0


Suggested change

supported_src_types = {ov.Type.u4: 4, ov.Type.i4: 4, ov.Type.u8: 8, ov.Type.i8: 8}

if ov_node.tensor_view.element_type not in supported_src_types:

return False, 0, 0

bits = supported_src_types[ov_node.tensor_view.element_type]

asym = ov_node.tensor_view.element_type in [ov.Type.u8, ov.Type.u4]

# only repack for sym types are supported for now: i3, i2

if asym:

return False, 0, 0

supported_src_types = {ov.Type.i4: 4, ov.Type.i8: 8}

if ov_node.tensor_view.element_type not in supported_src_types:

return False, 0, 0

?

AlexanderDokuchaev · 2026-09-29T11:24:43Z

+        bits = supported_src_types[ov_node.tensor_view.element_type]
+        asym = ov_node.tensor_view.element_type in [ov.Type.u8, ov.Type.u4]


Suggested change

bits = supported_src_types[ov_node.tensor_view.element_type]

asym = ov_node.tensor_view.element_type in [ov.Type.u8, ov.Type.u4]

bits = ov_node.tensor_view.element_type.get_bitwidth()

asym = ov_node.tensor_view.element_type.is_signed()

AlexanderDokuchaev · 2026-09-29T11:25:18Z

+
+        ov_node = self.name_to_node_mapping[node.node_name]
+
+        supported_src_types = {ov.Type.u4: 4, ov.Type.i4: 4, ov.Type.u8: 8, ov.Type.i8: 8}


Looks like now suppoerted only i4 and i8

AlexanderDokuchaev · 2026-09-29T12:25:07Z

+        if hasattr(weight_compression_parameters, "sequence"):
+            # avoid track
+            is_all_8bit = all(
+                wc_params.compression_config.num_bits == 8 for wc_params in weight_compression_parameters.sequence
+            )
+        else:
+            is_all_8bit = all(wc_params.compression_config.num_bits == 8 for wc_params in weight_compression_parameters)


Suggested change

if hasattr(weight_compression_parameters, "sequence"):

# avoid track

is_all_8bit = all(

wc_params.compression_config.num_bits == 8 for wc_params in weight_compression_parameters.sequence

)

else:

is_all_8bit = all(wc_params.compression_config.num_bits == 8 for wc_params in weight_compression_parameters)

is_all_8bit = all(

wc_params.compression_config.num_bits == 8

for wc_params in getattr(weight_compression_parameters, "sequence", weight_compression_parameters)

)

Looks like track uses alwasy here but ok

anzr299

Minor

anzr299 · 2026-09-29T13:06:36Z

@@ -0,0 +1,879 @@
+# Copyright (c) 2026 Intel Corporation


Maybe we can rename the example to QAD with Lora since it seems more suitable

anzr299 · 2026-09-29T13:08:18Z

+        transformers.enable_full_determinism(42)
+    else:
+        transformers.set_seed(42)


do we need the else statemnet? does transformers.set_seed(42) affect full_determinism()?

This parameter can slow down training, so I suggest keeping it optional.

…aining

Copilot

Copilot review overview

🟡 Changes recommended

Critical repacking and Torch backend issues, along with additional correctness and validation gaps, remain unresolved.

Review effort: Lite
Findings: 2 · 3 · 1

Open (6)

Select signed bit width based on the actual offset range · New
Reject or implement INT2/INT3 support for Torch DQ · New
Validate microbatch accumulation and flush partial batches · New
Preserve the original Convert element type when replacing it · New
Handle meta-device offloaded weights during equalization
Strengthen repacking test with graph/output validation

Resolved since last review (4)

Add 2/3-bit DQ packing support or reject validation
Identify zero-point constants by graph role, not shape
Include gate_proj in producer search
Use logger instead of unconditional print calls

+        max_val = max(int(weight_tensor.max().item()), abs(int(weight_tensor.min().item())))
+        optimal_bits = int(max_val).bit_length()
+
+        supported_optimal_bits = {2: ov.Type.u2, 3: ov.Type.u3}
+
+        # If the optimal bits is equal to the current bits, no need to repack
+        if optimal_bits == bits or optimal_bits not in supported_optimal_bits:


+            CompressWeightsMode.INT2_SYM: QuantizationScheme.SYMMETRIC_LORA,
+            CompressWeightsMode.INT3_SYM: QuantizationScheme.SYMMETRIC_LORA,


andreyanufr · 2026-10-01T16:26:35Z

+    opt = torch.optim.AdamW(param_to_train, weight_decay=weight_decay)
+
+    # Run tuning with distillation loss and validation after each epoch.
+    grad_accumulation_steps = args.batch_size // args.microbatch_size


andreyanufr · 2026-10-01T16:26:25Z

+        converted_const = opset.convert(compressed_const, ov.Type.f16)
+
+        zero_point_const = opset.constant(offset, dtype=ov.Type.i8, name=f"{const_node_name}/zero_point")
+        zero_point_const = opset.convert(zero_point_const, ov.Type.f16)
+
+        converted_const = opset.subtract(
+            converted_const, zero_point_const, name=f"{const_node_name}/zero_point/subtract"
+        )


2) Fixed convert type after repacking. 3) Added tokenizer to export.

…aining

MaximProshin · 2026-10-08T07:23:02Z

+
+This is a real training pipeline for low-bit LLM optimization, not a simple inference example.
+
+## Prerequisites


Is there any minimal or recommended HW? If so, let's mention it.

MaximProshin · 2026-10-08T07:27:57Z

+
+### Model and compression
+
+- `--pretrained`: HF model ID or local path.


I don't think it works with any model. Can we add a note somewhere which models this example was tested with?

MaximProshin · 2026-10-08T07:30:36Z

@@ -0,0 +1,123 @@
+# Quantization Aware Distillation with absorbable LoRA for 2/3-bit LLM compression


Let's try to explain all abbreviations when they're mentioned first time

andreyanufr added 3 commits July 27, 2026 17:39

Added rapcking of weights to support torch compression for 2-3 bit in…

32e2e5d

… nncf side.

1) Removed debug code.

029d4cf

2) Removed 6bit from repacking.

Fixed style bugs.

03ceee7

github-actions Bot added NNCF OpenVINO Pull requests that updates NNCF OpenVINO API Public API-impacting changes labels Jul 30, 2026

andreyanufr added 11 commits July 30, 2026 20:08

Added test for weight repacking.

db5a21a

Added test for weight repacking.

707c35f

Merge remote-tracking branch 'upstream/develop' into aanuf/2_3_bit_tr…

eeb5572

…aining

Train 2-3 bit quantizers.

2c45ac3

Revert experimental changes.

b0f0240

Fixed style and removed unused code.

ab544bb

1)Added optional equalization for MLP layer.

fc71994

2) Added optional possibility to save torch checkpoint with dequantized weights.

Avoid awq in the case of weight equalization.

8f8cc35

Added pile dataset for train and PTQ.

6ad2ffe

Add ultracaht dataset.

6d94fee

Added lr scheduler.

9fdee25

github-actions Bot added the documentation Improvements or additions to documentation label Sep 21, 2026

andreyanufr added 2 commits September 21, 2026 12:53

1) Removed lm_eval validation in the end of training.

f418fbb

2) Added README and requirements.txt

1) Added new words to spell check dict.

4bb71c7

2) Removed debug prints.

andreyanufr marked this pull request as ready for review September 24, 2026 14:57

Copilot AI lite review requested due to automatic review settings September 24, 2026 14:57

andreyanufr requested a review from a team as a code owner September 24, 2026 14:57

Copilot started reviewing on behalf of andreyanufr September 24, 2026 14:58 View session

Copilot AI reviewed Sep 24, 2026

View reviewed changes

andreyanufr added 2 commits September 24, 2026 17:24

Merge remote-tracking branch 'upstream/develop' into aanuf/2_3_bit_tr…

13a9cd4

…aining

1) Fixed logging of repack_weights.

547582c

2) Added arg for deterministic cuda execution. This arg is optional because it can slow down the training.

andreyanufr requested review from AlexanderDokuchaev and anzr299 September 25, 2026 10:53

AlexanderDokuchaev requested changes Sep 29, 2026

View reviewed changes

anzr299 approved these changes Sep 29, 2026

View reviewed changes

andreyanufr added 5 commits September 30, 2026 15:46

Rename QAT to QAD.

8ec441a

Fixed README.

0369755

Applied suggestions.

df4691d

Removed unused comment.

8c85da3

Merge remote-tracking branch 'upstream/develop' into aanuf/2_3_bit_tr…

f6e995b

…aining

andreyanufr requested a lite review from Copilot October 1, 2026 13:20

Copilot started reviewing on behalf of andreyanufr October 1, 2026 13:20 View session

Copilot AI reviewed Oct 1, 2026

View reviewed changes

andreyanufr added 10 commits October 1, 2026 15:38

Update optimum-intel version.

2a0130d

1) Added args check.

b414e89

2) Fixed convert type after repacking. 3) Added tokenizer to export.

Merge remote-tracking branch 'upstream/develop' into aanuf/2_3_bit_tr…

ebbeb18

…aining

Moved repack_weights to other place.

ecc6f5a

Moved repack_weights to other place.

8101998

Merge remote-tracking branch 'upstream/develop' into aanuf/2_3_bit_tr…

f40ee23

…aining

Added right scale merging for case of Qwen3NextRMSNorm.

e2df895

Changes imports order.

ffc7fa6

Disable autocast of bfloat16 matmul output to fp32.

1f558e0

Expand layer norm types for (1+w).

fc8c8f2

MaximProshin reviewed Oct 8, 2026

View reviewed changes

github-actions Bot removed the API Public API-impacting changes label Oct 8, 2026

andreyanufr added 2 commits October 8, 2026 14:28

Move import of ov module to function to avoid error in tests.

02005cf

Added type ignore.

502ecaa

		cd examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit
		python main.py --pretrained meta-llama/Llama-3.2-1B-Instruct --bits 3 --dataset pile-10k --epochs 1 --output_dir output

	cd examples/llm_compression/torch/distillation_qat_with_lora_2_3_bit
	python main.py --pretrained meta-llama/Llama-3.2-1B-Instruct --bits 3 --dataset pile-10k --epochs 1 --output_dir output
	python main.py

	- `--microbatch_size`: per-step microbatch size
	- `--microbatch_size`: per-step microbatch size
	- `--full_determinism`: to enable deterministic QAT

		bits = supported_src_types[ov_node.tensor_view.element_type]
		asym = ov_node.tensor_view.element_type in [ov.Type.u8, ov.Type.u4]


		ov_node = self.name_to_node_mapping[node.node_name]

		supported_src_types = {ov.Type.u4: 4, ov.Type.i4: 4, ov.Type.u8: 8, ov.Type.i8: 8}

		CompressWeightsMode.INT2_SYM: QuantizationScheme.SYMMETRIC_LORA,
		CompressWeightsMode.INT3_SYM: QuantizationScheme.SYMMETRIC_LORA,


		This is a real training pipeline for low-bit LLM optimization, not a simple inference example.

		## Prerequisites


		### Model and compression

		- `--pretrained`: HF model ID or local path.

		@@ -0,0 +1,123 @@
		# Quantization Aware Distillation with absorbable LoRA for 2/3-bit LLM compression

Conversation

andreyanufr commented Jul 30, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Reason for changes

Related tickets

Tests

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anzr299 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaximProshin Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

andreyanufr commented Jul 30, 2026 •

edited by github-actions Bot

Loading

MaximProshin Oct 8, 2026 •

edited

Loading